library(Seurat)
library(tidyverse)
library(tidyseurat)
source("00_util_scripts/mod_seurat.R", echo = FALSE)
source("DE_cells/scripts/DE_filtering_funcs.R")

monaco <- celldex::MonacoImmuneData()

disease_list <- c("ms", "sle", "pss")

# MS data ------
# https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi?acc=GSE138266
# wget -r ftp://ftp.ncbi.nlm.nih.gov/geo/series/GSE138nnn/GSE138266/suppl/
for (disease in disease_list) {
  if (disease == "ms") {
    ms_path <- c("matrix", "barcode", "gene") %>%
      map(\(x)list.files("DE_cells/data/MS/", pattern = x, full.names = TRUE))

    ms_sample_name <- ms_path[[1]] %>%
      map(str_extract, "GSM.+(CSF|PBMC)")

    sobj <- ms_path %>%
      pmap(ReadMtx, .progress = TRUE) %>%
      map2(ms_sample_name, add_name_field) %>%
      map(CreateSeuratObject, min.cells = 3, min.features = 200, names.field = 2, .progress = TRUE) %>%
      purrr::reduce(merge)
    
    sobj <- sobj %>%
      AddMetaData(colnames(sobj), "colname") %>%
      separate(col = colname, into = c("barcode", "batch", "group", "tissue"), sep = "_", remove = FALSE) |>
      unite("tissue_group", group:tissue, remove = FALSE) |>
      mutate(group = str_extract(group, "[a-zA-Z]+"), orig.ident = tissue_group)

    Idents(sobj) <- "orig.ident"
  }

  # SLE data ------
  # GSE162577
  # file just downloaded from geo is rar in nature, not gz
  # for file in *.gz ; do unrar e ${file} ; done
  if (disease == "sle") {
    sle_path <- c("matrix", "barcode", "gene") %>%
      map(\(x)list.files("DE_cells/data/SLE/", pattern = x, full.names = TRUE))

    sle_sample_name <- sle_path[[1]] %>%
      map(\(x)str_extract(x, "(SLE|C)-(1|2)"))

    sobj <- sle_path %>%
      pmap(\(x, y, z)ReadMtx(mtx = x, cells = y, features = z), .progress = TRUE) %>%
      map2(sle_sample_name, add_name_field) %>%
      map(\(x)CreateSeuratObject(x, min.cells = 3, min.features = 200, names.field = 2), .progress = TRUE) %>%
      purrr::reduce(merge)
  }

  # pSS data -------
  # wget -r ftp://ftp.ncbi.nlm.nih.gov/geo/series/GSE157nnn/GSE157278/suppl/'
  if (disease == "pss") {
    pss_path <- c("matrix", "barcode", "features") %>%
      map(\(x)list.files("DE_cells/data/pSS/", pattern = x, full.names = TRUE))

    sobj <- ReadMtx(mtx = pss_path[[1]], cells = pss_path[[2]], features = pss_path[[3]]) %>%
      CreateSeuratObject(min.cells = 3, min.features = 200)

    pss_meta <- read_tsv("DE_cells/data/pSS/GSE157278_cell_batch.tsv.gz")

    sobj <- sobj %>%
      left_join(pss_meta, by = c(".cell" = "Cell")) %>%
      mutate(orig.ident = batch)
    
    Idents(sobj) <- 'orig.ident'
  }
  
  # read perez SLE data ----------
  sobj <- read_rds('DE_cells/data/perez_sle.rds')

  # process seurat ---------
  sobj$mitoRatio <- PercentageFeatureSet(object = sobj, pattern = "^MT-")
  
  sobj$orig.ident <- sobj$batch_cov

  VlnPlot(sobj, "mitoRatio", pt.size = 0)

  sobj <- sobj %>%
    filter(mitoRatio < 10)

  sobj <- sobj %>%
    mark_doublets() |>
    filter(scrublet_call == "Singlet")
  
  sobj <- sobj %>%
    quick_process_seurat()
  
  sobj <- NormalizeData(sobj)
  
  sobj <- sobj |>
    FindVariableFeatures() |>
    ScaleData() |>
    RunPCA() 
  
  sobj <- sobj |>
    RunHarmony() |>
    RunUMAP(reduction = "harmony", dims = 1:20) |>
    FindNeighbors(reduction = "harmony", dims = 1:20) 

  sobj <-
    sobj |>
    FindClusters()
  
  DimPlot(sobj, label = TRUE, label.box = TRUE) |> print()

  DimPlot(sobj, group.by = 'Status')
    
  sobj_meta <- sobj %>%
    as_tibble() %>%
    select(!matches("PC_|harmony")) %>%
    write_csv(str_glue("DE_cells/results/{disease}_sobj_meta.csv"))

  # adjust by soupx --------
  # soupx need clustering info to run
  souped_sobj <- sobj |>
    DietSeurat() |>
    SplitObject('orig.ident') |>
    map(soupx_seurat, .progress = TRUE) |>
    purrr::reduce(merge)
  
  souped_sobj |>
    write_rds(str_glue("DE_cells/data/{disease}_adj_sobj.rds"))

  perez_sce <- zellkonverter::readH5AD('~/learn/scvelo/perez_sle.h5ad', X_name = 'counts',
                          reader = 'R',
                          verbose = TRUE)
  
  perez_sce |>
    mutate(new_cb = str_remove(.cell, '-.+') |>
             str_c(ind_cov, batch_cov, cg_cov, sep = '_')) |>
    pull(new_cb) -> new_cb
  
  colnames(perez_sce) <- new_cb |>
    make.names(unique = TRUE)
  
  colData(perez_sce) -> perez_meta
  
  sobj_perez <- assay(perez_sce) |>
    CreateSeuratObject(min.cells = 3, min.features = 200)
  
  # cell type annotation ------
  sobj <- sobj |>
    mark_cell_type_singler(
      ref = monaco,
      fine_label = TRUE,
      new_label = "monaco_label"
    )
  
  sobj <- sobj |>
    mark_cell_type_singler(
      ref = monaco,
      fine_label = FALSE,
      new_label = "monaco_major"
    )
  
  sobj <- sobj |>
    mark_cell_type_singler(
      ref = dice_ref,
      fine_label = TRUE,
      new_label = "dice_label"
    )
  
  DimPlot(sobj, group.by = 'monaco_label', label = TRUE, repel = TRUE)
  
  Idents(sobj) <- 'monaco_label'
  
  # plot cluster markers
  pbmc_markers <- FindAllMarkers(sobj, only.pos = TRUE, logfc.threshold = 1, min.pct = 0.3)
  
  best_markers <- pbmc_markers |>
    group_by(cluster) |>
    slice_max(avg_log2FC, n = 2) |>
    pull(gene) |> unique()
  
  sobj |> DotPlot(features = best_markers) +
    theme(axis.text.x = element_text(angle = 90, hjust = 1, vjust = .5))

  write_rds(sobj, str_glue("DE_cells/data/{disease}_sobj.rds"))

  # annotate DE cell features -----
  get_abundance_sc_wide(sobj, all_marker) %>%
    right_join(as_tibble(sobj)) %>%
    select(!matches("PC_|harmony")) %>%
    data.table::fwrite(str_glue("DE_cells/results/{disease}_de_meta.csv.gz"))
}

